//
//  FQDMultiColumnsCollectionViewLayout.m
//  FakeQDaily
//
//  Created by mewxiv on 9/14/16.
//  Copyright © 2016 Xensius. All rights reserved.
//

#import "FQDMultiColumnsCollectionViewLayout.h"

static NSInteger const MaxUnionSize = 20;

@interface FQDMultiColumnsCollectionViewLayout ()

@property (nonatomic, strong) NSMutableArray *columnHeights;
@property (nonatomic, strong) NSMutableArray *columnWidths;
@property (nonatomic, strong) NSArray *sectionItemAttributes;
@property (nonatomic, strong) NSArray *allItemAttributes;
@property (nonatomic, strong) NSArray *unionRects;

@end

@implementation FQDMultiColumnsCollectionViewLayout

- (instancetype)init
{
    self = [super init];
    if (self) {
        _columnCount = 2;
        _minimumColumnSpacing = 10.0f;
        _minimumLineSpacing = 10.0f;
    }
    return self;
}

- (void)prepareLayout
{
    [super prepareLayout];
    
    NSInteger sectionCount = [self.collectionView numberOfSections];
    if (sectionCount == 0) {
        return;
    }
    
    NSMutableArray *muSectionItemAttributes = [NSMutableArray array];
    NSMutableArray *muAllItemAttributes = [NSMutableArray array];
    NSMutableArray *muUnionRects = [NSMutableArray array];
    self.columnWidths = [NSMutableArray array];
    self.columnHeights = [NSMutableArray array];
    
    CGFloat offsetY = 0;
    UICollectionView *collectionView = self.collectionView;
    
    for (NSUInteger section = 0; section < sectionCount; section++) {
        CGFloat minimumLineSpacing = self.minimumLineSpacing;
        if ([self.delegate respondsToSelector:@selector(collectionView:layout:minimumLineSpacingForSectionAtIndex:)]) {
            minimumLineSpacing = [self.delegate collectionView:collectionView layout:self minimumLineSpacingForSectionAtIndex:section];
        }
        
        UIEdgeInsets sectionInset = self.sectionInset;
        if ([self.delegate respondsToSelector:@selector(collectionView:layout:insetForSectionAtIndex:)]) {
            sectionInset = [self.delegate collectionView:collectionView layout:self insetForSectionAtIndex:section];
        }
        
        NSInteger columnCount = self.columnCount;
        if ([self.delegate respondsToSelector:@selector(collectionView:columnCountForSection:)]) {
            columnCount = [self.delegate collectionView:self.collectionView columnCountForSection:section];
        }
        
        CGFloat availableWidth = (collectionView.frame.size.width - sectionInset.left - sectionInset.right);
        CGFloat itemWidth = floorf((availableWidth + self.minimumColumnSpacing) / columnCount - self.minimumColumnSpacing);
        [self.columnWidths addObject:@(itemWidth)];
        
        offsetY += sectionInset.top;
        for (NSUInteger column = 0; column < columnCount; column++) {
            self.columnHeights[column] = @(offsetY);
        }
        
        NSUInteger itemCount = [collectionView numberOfItemsInSection:section];
        NSMutableArray *muItemAttributes = [NSMutableArray array];
        
        for (NSUInteger i = 0; i < itemCount; i++) {
            NSIndexPath *indexPath = [NSIndexPath indexPathForItem:i inSection:section];
            NSUInteger columnIndex = [self shortestColumnIndex];
            
            CGFloat x = sectionInset.left + (itemWidth + self.minimumColumnSpacing) * columnIndex;
            CGFloat y = [self.columnHeights[columnIndex] floatValue];
            CGSize itemSize = [self.delegate collectionView:collectionView layout:self sizeForItemAtIndexPath:indexPath];
            CGFloat itemHeight = 0;
            if (itemSize.width > 0 && itemSize.height > 0) {
                itemHeight = itemSize.height * itemWidth / itemSize.width;
            }
            
            UICollectionViewLayoutAttributes *attributes = [UICollectionViewLayoutAttributes layoutAttributesForCellWithIndexPath:indexPath];
            attributes.frame = CGRectMake(x, y, itemWidth, itemHeight);
            [muItemAttributes addObject:attributes];
            [muAllItemAttributes addObject:attributes];
            self.columnHeights[columnIndex] = @(CGRectGetMaxY(attributes.frame) + minimumLineSpacing);
        }
        
        [muSectionItemAttributes addObject:muItemAttributes];
        
        NSUInteger columnIndex = [self tallestColumnIndex];
        offsetY = [self.columnHeights[columnIndex] floatValue] - minimumLineSpacing + sectionInset.bottom;
        
        for (NSUInteger column = 0; column < columnCount; column++) {
            self.columnHeights[column] = @(offsetY);
        }
    }
    
    self.sectionItemAttributes = [muSectionItemAttributes copy];
    self.allItemAttributes = [muAllItemAttributes copy];
    
    NSUInteger itemIndex = 0;
    NSUInteger itemCount = [self.allItemAttributes count];
    while (itemIndex < itemCount) {
        UICollectionViewLayoutAttributes *beginAttributes = self.allItemAttributes[itemIndex];
        itemIndex = MIN(itemIndex + MaxUnionSize, itemCount) - 1;
        UICollectionViewLayoutAttributes *endAttributes = self.allItemAttributes[itemIndex];
        CGRect beginRect = beginAttributes.frame;
        CGRect endRect = endAttributes.frame;
        [muUnionRects addObject:[NSValue valueWithCGRect:CGRectUnion(beginRect, endRect)]];
        itemIndex++;
    }
    self.unionRects = [muUnionRects copy];
}

- (CGSize)collectionViewContentSize
{
    NSUInteger sectionCount = [self.collectionView numberOfSections];
    if (sectionCount == 0) {
        return CGSizeZero;
    }
    
    CGSize contentSize = self.collectionView.bounds.size;
    contentSize.height = [self.columnHeights[0] floatValue];
    
    return contentSize;
}

- (UICollectionViewLayoutAttributes *)layoutAttributesForItemAtIndexPath:(NSIndexPath *)indexPath
{
    return self.sectionItemAttributes[indexPath.section][indexPath.row];
}

- (NSArray<UICollectionViewLayoutAttributes *> *)layoutAttributesForElementsInRect:(CGRect)rect
{
    NSUInteger begin = 0;
    NSUInteger end = [self.unionRects count];
    NSMutableArray *attrsList = [NSMutableArray array];
    
    for (NSInteger i = 0; i < [self.unionRects count]; i++) {
        if (CGRectIntersectsRect(rect, [self.unionRects[i] CGRectValue])) {
            begin = i * MaxUnionSize;
            break;
        }
    }
    
    for (NSInteger i = [self.unionRects count]-1; i >= 0; i--) {
        if (CGRectIntersectsRect(rect, [self.unionRects[i] CGRectValue])) {
            end = MIN((i + 1) * MaxUnionSize, [self.allItemAttributes count]);
            break;
        }
    }
    
    for (NSInteger i = begin; i < end; i++) {
        UICollectionViewLayoutAttributes *attrs = self.allItemAttributes[i];
        attrs.zIndex = end - i;
        if (CGRectIntersectsRect(rect, attrs.frame)) {
            [attrsList addObject:attrs];
        }
    }
    
    return [attrsList copy];
}

- (BOOL)shouldInvalidateLayoutForBoundsChange:(CGRect)newBounds
{
    CGRect oldBounds = self.collectionView.bounds;
    if (CGRectGetWidth(oldBounds) != CGRectGetWidth(newBounds)) {
        return YES;
    }
    
    return NO;
}

- (NSUInteger)shortestColumnIndex
{
    NSUInteger index = 0;
    CGFloat shortestHeight = CGFLOAT_MAX;
    
    for (NSNumber *heightNumber in self.columnHeights) {
        CGFloat height = [heightNumber floatValue];
        if (height < shortestHeight) {
            shortestHeight = height;
            index = [self.columnHeights indexOfObject:heightNumber];
        }
    }
    
    return index;
}

- (NSUInteger)tallestColumnIndex
{
    NSUInteger index = 0;
    CGFloat tallestHeight = 0;
    
    for (NSNumber *heightNumber in self.columnHeights) {
        CGFloat height = [heightNumber floatValue];
        if (height > tallestHeight) {
            tallestHeight = height;
            index = [self.columnHeights indexOfObject:heightNumber];
        }
    }
    
    return index;
}

- (CGFloat)columnWidthForSection:(NSUInteger)section
{
    return [self.columnWidths[section] floatValue];
}

- (id<FQDMultiColumnsCollectionViewLayoutDelegate>)delegate
{
    return (id<FQDMultiColumnsCollectionViewLayoutDelegate>)self.collectionView.delegate;
}

@end
